package io.kiki.sba.registry.server.clientInterface.store;


import io.kiki.sba.registry.common.model.Tuple;

import java.util.*;
import java.util.concurrent.ConcurrentHashMap;
import java.util.function.BiConsumer;

public class SlotStore<T> implements Store<T> {
    private final Map<Integer, _store_<T>> map = new ConcurrentHashMap<>(512);
    private final SlotOfCaller slotOfCaller;

    public SlotStore(SlotOfCaller slotOfCaller) {
        this.slotOfCaller = slotOfCaller;
    }


    private _store_<T> getOrCreateStore(String dataInfoId) {
        int slotId = slotOfCaller.slotOf(dataInfoId);
        return getOrCreateStore(slotId);
    }

    private _store_<T> getOrCreateStore(int slotId) {
        return map.computeIfAbsent(slotId, k -> new _store_<>(256, 256));
    }

    @Override
    public Map<String, T> get(String dataInfoId) {
        return getOrCreateStore(dataInfoId).get(dataInfoId);
    }

    @Override
    public Map<String, T> getOrCreate(String dataInfoId) {
        return getOrCreateStore(dataInfoId).getOrCreate(dataInfoId);
    }

    @Override
    public void forEach(BiConsumer<String, Map<String, T>> consumer) {
        for (_store_<T> store : map.values()) {
            store.forEach(consumer);
        }
    }

    @Override
    public Map<String, Map<String, T>> copyMap() {
        Map<String, Map<String, T>> map1 = new HashMap<>(1024 * 16);
        for (_store_<T> store : map.values()) {
            map1.putAll(store.copyMap());
        }
        return map1;
    }

    public Map<String, Map<String, T>> copyMap(int slotId) {
        _store_<T> store = getOrCreateStore(slotId);
        return store.copyMap();
    }

    public Tuple<Long, Long> count() {
        long dataInfoIdCount = 0;
        long dataCount = 0;
        for (_store_<T> store : map.values()) {
            Tuple<Long, Long> tuple = store.count();
            dataInfoIdCount += tuple.getFirst();
            dataCount += tuple.getSecond();
        }
        return new Tuple<>(dataInfoIdCount, dataCount);
    }

    @Override
    public Collection<String> getDataInfoIds() {
        Set<String> set = new HashSet<>(1024 * 16);
        for (_store_<T> store : map.values()) {
            set.addAll(store.getDataInfoIds());
        }
        return set;
    }

    public interface SlotOfCaller {
        int slotOf(String dataInfoId);
    }
}
